import cv2
import numpy as np
import matplotlib.pyplot as plt
import os

def fft2_r(pic_path):
    """
    对图像傅里叶变换，再根据亮点进行逆变换
    """
    # 1. 读图
    img = cv2.imread(pic_path, cv2.IMREAD_GRAYSCALE).astype(np.float32)

    # 2. FFT + 中心化
    fft = np.fft.fft2(img)
    fft_shift = np.fft.fftshift(fft)  # 中心化后的复数谱
    mag = np.abs(fft_shift)  # 幅度谱

    # 3. 阈值找亮点（幅度 > 90% 最大值）
    thr = np.percentile(mag, 99)  # 也可手动设阈值
    mask = (mag >= thr)

    # 4. 强制共轭对称：若某点保留，其共轭对称点也必须保留
    h, w = mask.shape
    for y in range(h):
        for x in range(w):
            if mask[y, x]:
                # 共轭对称位置
                y2, x2 = h - 1 - y, w - 1 - x
                mask[y2, x2] = True

    # 5. 应用掩膜
    fft_shift_masked = fft_shift * mask

    # 6. 逆变换
    recon = np.fft.ifft2(np.fft.ifftshift(fft_shift_masked)).real

    # 7. 可视化
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 3, 1)
    plt.title('Original')
    plt.imshow(img, cmap='gray')
    plt.subplot(1, 3, 2)
    plt.title('Center-symmetric Bright Spots')
    plt.imshow(mask, cmap='gray')
    plt.subplot(1, 3, 3)
    plt.title('Reconstructed')
    plt.imshow(np.clip(recon, 0, 255), cmap='gray')
    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    pic_path = os.path.join('../img', 'new', '0007.tif')
    fft2_r(pic_path)